/*
A demo PE-loader for the training: https://github.com/hasherezade/malware_training_vol1
WARNING: This is a basic example. For the sake of simplicity, some checks have been omitted.
*/

#include <windows.h>
#include <iostream>

#define RELOC_32BIT_FIELD 3
#define RELOC_64BIT_FIELD 0xA

#ifdef _WIN64
#define RELOC_FIELD RELOC_64BIT_FIELD
typedef ULONG_PTR FIELD_PTR;
#else
#define RELOC_FIELD RELOC_32BIT_FIELD
typedef  DWORD_PTR FIELD_PTR;
#endif

typedef struct _BASE_RELOCATION_ENTRY {
    WORD Offset : 12;
    WORD Type : 4;
} BASE_RELOCATION_ENTRY;

inline void manual_map(BYTE* image, BYTE *rawPE, PIMAGE_NT_HEADERS nt)
{
    memcpy(image, rawPE, nt->OptionalHeader.SizeOfHeaders);

    // map sections
    PIMAGE_SECTION_HEADER section = IMAGE_FIRST_SECTION(nt);
    for (WORD i = 0; i < nt->FileHeader.NumberOfSections; i++) {

        memcpy((BYTE*)(image)+section[i].VirtualAddress, (BYTE*)(rawPE)+section[i].PointerToRawData, section[i].SizeOfRawData);
    }
}

inline bool relocate(BYTE* image, PIMAGE_NT_HEADERS nt, FIELD_PTR newImgBase)
{
    IMAGE_DATA_DIRECTORY relocationsDirectory = nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC];
    if (relocationsDirectory.VirtualAddress == 0) {
        return false;
    }
    PIMAGE_BASE_RELOCATION ProcessBReloc = (PIMAGE_BASE_RELOCATION)(relocationsDirectory.VirtualAddress + (FIELD_PTR)image);
    // apply relocations:
    while (ProcessBReloc->VirtualAddress != 0)
    {
        DWORD page = ProcessBReloc->VirtualAddress;
#ifdef _DEBUG
        std::cout << "page: " << std::hex << page << std::endl;
#endif
        if (ProcessBReloc->SizeOfBlock >= sizeof(IMAGE_BASE_RELOCATION))
        {
            size_t count = (ProcessBReloc->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(WORD);
            BASE_RELOCATION_ENTRY* list = (BASE_RELOCATION_ENTRY*)(LPWORD)(ProcessBReloc + 1);
#ifdef _DEBUG
            std::cout << "Count: " << count << ":\n";
#endif
            for (size_t i = 0; i < count; i++)
            {
                if (list[i].Type & RELOC_FIELD)
                {
                    DWORD rva = list[i].Offset + page;
#ifdef _DEBUG
                    std::cout << "RVA : " << std::hex << rva << "\n";
#endif
                    PULONG_PTR p = (PULONG_PTR)((LPBYTE)image + rva);
                    //relocate the address
                    *p = ((*p) - nt->OptionalHeader.ImageBase) + (FIELD_PTR)newImgBase;
                }
            }
        }
#ifdef _DEBUG
        std::cout << "---\n";
#endif
        ProcessBReloc = (PIMAGE_BASE_RELOCATION)((LPBYTE)ProcessBReloc + ProcessBReloc->SizeOfBlock);
    }
    return true;
}

inline bool load_imports(BYTE* image, PIMAGE_NT_HEADERS nt)
{
    IMAGE_DATA_DIRECTORY importsDirectory = nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
    if (importsDirectory.VirtualAddress == 0) {
        return false;
    }
    PIMAGE_IMPORT_DESCRIPTOR importDescriptor = (PIMAGE_IMPORT_DESCRIPTOR)(importsDirectory.VirtualAddress + (FIELD_PTR)image);

    while (importDescriptor->Name != NULL)
    {
        LPCSTR libraryName = (LPCSTR)importDescriptor->Name + (FIELD_PTR)image;
        HMODULE library = LoadLibraryA(libraryName);
#ifdef _DEBUG
        std::cout << "Loading: " << libraryName << "\n";
#endif
        if (library)
        {
            PIMAGE_THUNK_DATA thunk = NULL;
            thunk = (PIMAGE_THUNK_DATA)((FIELD_PTR)image + importDescriptor->FirstThunk);

            while (thunk->u1.AddressOfData != NULL)
            {
                FIELD_PTR functionAddress = NULL;
                if (IMAGE_SNAP_BY_ORDINAL(thunk->u1.Ordinal))
                {
                    LPCSTR functionOrdinal = (LPCSTR)IMAGE_ORDINAL(thunk->u1.Ordinal);
                    functionAddress = (FIELD_PTR)GetProcAddress(library, functionOrdinal);
                }
                else
                {
                    PIMAGE_IMPORT_BY_NAME functionName = (PIMAGE_IMPORT_BY_NAME)((FIELD_PTR)image + thunk->u1.AddressOfData);
                    functionAddress = (FIELD_PTR)GetProcAddress(library, functionName->Name);
                }
                thunk->u1.Function = functionAddress;
                ++thunk;
            }
        }

        importDescriptor++;
    }
    return true;
}

PIMAGE_NT_HEADERS get_nt_hdr(BYTE *rawPE)
{
    //get header
    IMAGE_DOS_HEADER* DOSHeader = PIMAGE_DOS_HEADER(rawPE);
    if (DOSHeader->e_magic != IMAGE_DOS_SIGNATURE) {
        return NULL;
    }
    PIMAGE_NT_HEADERS nt = PIMAGE_NT_HEADERS((char*)(rawPE)+DOSHeader->e_lfanew);
    if (nt->Signature != IMAGE_NT_SIGNATURE) {
        return NULL;
    }
    return nt;
}

bool load_and_run(BYTE* rawPE, size_t rawSize)
{
    //get header
    IMAGE_DOS_HEADER* DOSHeader = PIMAGE_DOS_HEADER(rawPE);
    PIMAGE_NT_HEADERS nt = PIMAGE_NT_HEADERS((char*)(rawPE)+DOSHeader->e_lfanew);
    if (!nt) {
        std::cerr << "Not a PE file\n";
        return false;
    }
    //write file to memory
    BYTE* image = (BYTE*)VirtualAlloc(NULL, nt->OptionalHeader.SizeOfImage, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
    if (!image) {
        std::cerr << "Allocating image has failed\n";
        return false;
    }
#ifdef _DEBUG
    std::cout << "Allocated at:" << std::hex << image << "\n";
#endif
    manual_map(image, rawPE, nt);
    if (!relocate(image, nt, (FIELD_PTR)image)) {
        std::cerr << "Relocating image has failed\n";
        return false;
    }
    load_imports(image, nt);

    // call new Entry Point
    ULONG_PTR EntryPoint = nt->OptionalHeader.AddressOfEntryPoint + (ULONG_PTR)image;
#ifdef NEW_THREAD
    HANDLE hThread = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)EntryPoint, NULL, 0, NULL);
    if (hThread) {
        std::cout << "Thread run at: " << std::hex << EntryPoint << "\n";
        WaitForSingleObject(hThread, INFINITE);
    }
#else
    std::cout << "Calling Entry Point at: " << std::hex << EntryPoint << "\n";
    int (*new_main)() = (int(*)())EntryPoint;
    //call the Entry Point of the manually loaded PE:
    new_main();
#endif
    return true;
}

BYTE* read_file(IN const LPSTR sourcePath, OUT DWORD &loaded_size)
{
    HANDLE hFile = CreateFileA(sourcePath, GENERIC_READ, 0, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (!hFile || hFile == INVALID_HANDLE_VALUE) {
        std::cerr << "Opening the file: " << sourcePath << " has failed! Error: " << std::dec << GetLastError() << "\n";
        return NULL;
    }
    DWORD size = GetFileSize(hFile, NULL);
    if (size == INVALID_FILE_SIZE) {
        CloseHandle(hFile);
        return NULL;
    }
    //read the file
    BYTE* rawPE = new BYTE[size];
    if (!ReadFile(hFile, rawPE, size, NULL, NULL)) {
        std::cerr << "[ERROR] Reading the file has failed!\n";
        delete[]rawPE;
        rawPE = NULL;
    }
    CloseHandle(hFile);
    return rawPE;
}

int main(int argc, char* argv[])
{
    //get source process
    if (argc < 2) {
        std::cout << "pe_self_load: ";
        std::cout << "a manual PE loader with self-injection\n";
#ifdef _WIN64
        std::cout << "64-bit version" << "\n";
#else
        std::cout << "32-bit version" << "\n";
#endif
        std::cout << "Args: <exe_path>" << "\n";
        return 1;
    }

    const LPSTR sourcePath = argv[1];
    DWORD size = 0;
    BYTE* rawPE = read_file(sourcePath, size);
    if (!rawPE) {
        return 0;
    }
    if (!load_and_run(rawPE, size)) {
        std::cerr << "[ERROR] Loading failed!\n";
    }
    return 0;
}
